import json
from src.agents.file_qa import file_qa
from src.huey_init._huey import huey
from aimodel import tongyi_model
from src.tools.dms_operator import DMSOperator
from src.tools.extract_jsonstring import re_json_string_list
from fastapi import APIRouter
from pydantic import BaseModel

def find_body(index, paragraphInfos):
    for item in paragraphInfos:
        if item["index"] == index:
            return item

@huey.task(context=True)
def correct_task(correct, task=None):
    for file_ids_json in correct.exchange_info["docList"]:
        file_id = file_ids_json["realPath"]
        dms_op =DMSOperator()
        file_json_body = dms_op.get_jsonfile(file_id)
        pageCount = file_json_body.get("pageCount")
        paragraphInfos = file_json_body.get("paragraphInfos")
        tableInfos = file_json_body.get("tableInfos")
        # print(file_json_body)
        edit_file = {
            "fileId":file_id,
            "editContentInfos":[]
        }
        full_context = []
        for p in paragraphInfos:
            full_context.append({"index":p["index"],"text":p["text"]})
        example = [{"index":0, "before":"纠正前段落内容", "after":"纠正后段落内容"}]
        full_context_str = json.dumps(full_context,ensure_ascii=False)
        # prompt = ("你是一个文件审查员，请审查以下文件内容；每个对象表示一个文字段落；index表示段落编号，text表示段落内容\n"
        #           f"```json\n{full_context_str}\n```\n"
        #           f"请找出文中段落中可能会出现的语法或者名词写法错误；用以下格式回答\n```json\n{example}\n```\n"
        #           "只用按json格式回答，不要有多余的话。"#如果没有都；回答: “[]” ，
        #           )
        prompt = ("你是一位专业的文件审核员，请您检查并指出下列文档中可能存在的语法错误或错字或名词错误。"
                  "以下内容中每个对象代表一个独立的段落，其中“index”为段落编号，“text”为段落的具体内容。"
                  f"```json\n{full_context_str}\n```\n"
                  f"请您按照以下示例格式给出有错误的部分信息，\n```json\n{example}\n```\n"
                  f"如果没有找到错误，回答\n```\njson\n[]\n```\n只用给出有错误的段落，请仅以JSON格式回复，避免任何额外的说明。"
                  )
        # print(prompt)
        stream_res = tongyi_model.stream(prompt)
        result = ""
        for a in stream_res:
            print(a.content,end="")
            result = result+a.content
        res_str = re_json_string_list(result)
        try:
            res_json = json.loads(res_str)
        except:
            res_json = eval(res_str)
        for i in res_json:
            edit_file["editContentInfos"].append(
                {"index":i["index"],
                 "sectionIndex":find_body(i["index"],paragraphInfos)["sectionIndex"],#
                 "bodyChildIndex":find_body(i["index"],paragraphInfos)["bodyChildIndex"],
                 "operationType":0,
                 "documentObjectType":0,
                 "content": i["after"].strip()
                 }
            )
        res_edit = dms_op.post_edit_file(edit_file)
        # 调用修改对比接口
        res2 = dms_op.post_compair(file_id,res_edit["data"])
        res_path = res2["data"]
        # dms_op.post_exchange_info()
        exchange_info = {
            "functionType": correct.exchange_info.get("functionType"),
            "taskType": correct.exchange_info.get("taskType"),
            "taskDocumentType": correct.exchange_info.get("taskDocumentType"),
            "executionType": correct.exchange_info.get("executionType"),
            "isSuccess": True,
            "dmsTaskId": correct.exchange_info.get("dmsTaskId"),
            "docId": file_ids_json.get("docId"),
            "taskDocId": file_ids_json.get("taskDocId"),
            "resultFirstFilePath": res_path,
            "resultSecondFilePath": None,
            "aiTaskId": task.id,
            "aiGenFileType": 0,
        }
        dms_op.post_exchange_info(exchange_info)


class Corrrect(BaseModel):
    error_correct_file_list:list
    correct_indicator_list:list
    need_consistency_check:bool
    execution_type: int
    exchange_info:dict


router = APIRouter(
    prefix="/correct",
    tags=["correct"],
    responses={404: {"description": "Not found"}}
)

@router.post("/correct")
async def correct_fnc(correct: Corrrect):
    res = correct_task(correct)
    return {"status": "success", "aitask_id": res.id}


if __name__ == "__main__":
    pass
    #correct_task("Files\\202411\\20241108142016eloEldpFpkB线葡萄糖注射液原版2-纠错测试.docx")